#!/usr/bin/env python3
#coding=utf8

import time
import uuid
import hashlib
import peewee

USERDATA = './data/userdata.db'
userdata = peewee.SqliteDatabase(USERDATA, threadlocals = True)
userdata.connect()

def decorator(f):
    def F1(*args,**kwargs):
        result = f(*args,**kwargs)
        return result
    return F1

class BaseModel(peewee.Model):
    class Meta:
        database = userdata

###########################################################################################
class User(BaseModel):
    id           = peewee.PrimaryKeyField()
    umobile      = peewee.CharField(unique = True, max_length = 11, null = False)    #注册手机号
    upwd         = peewee.CharField(max_length = 32, null = False)                   #密码
    uname        = peewee.CharField(max_length = 20, default='')                     #性名
    uemail       = peewee.CharField(max_length = 20, default='')                     #邮箱
    ugender      = peewee.IntegerField(default = 0)                                  #性别 1-男 2-女 0-其它
    utype        = peewee.IntegerField(default = 0)                                  #用户类型 0-user 1-admin
    createdate   = peewee.IntegerField(default = int(time.time()))                   #创建时间
    lastlogin    = peewee.IntegerField(default = int(time.time()))                   #最后登录时间

    @staticmethod
    def create_password(raw,sale='\r\a\b\c\d\e\n'):
        return hashlib.new("md5", (sale ^ raw.encode())).hexdigest()[::-1]

    def check_password(self, raw):
        return self.create_password(raw) == self.upwd

    def update_login(self):
        self.lastlogin = int(time.time())
        self.save()

######################################################################################
class Plugin(BaseModel):
    """插件表"""
    plugin_type  = peewee.CharField(unique = True)
    plugin_path  = peewee.CharField(default = "portscan.py")
    plugin_setup = peewee.BooleanField(default = True)
    plugin_name  = peewee.CharField(default = "IP端口扫描")
    createdate   = peewee.IntegerField(default = int(time.time()))

class Task(BaseModel):
    """扫描任务表"""
    id           = peewee.PrimaryKeyField()
    task_id      = peewee.CharField(unique = True, default = uuid.uuid1())
    task_pid     = peewee.IntegerField(default = 0)   #进程PID
    task_state   = peewee.IntegerField(default = 0)                 #0-未开始 1-运行中 2-已完成 3-出错
    task_level   = peewee.IntegerField(default = 0)                 #任务等级 0、1、2
    starttime    = peewee.IntegerField(default = int(time.time()))  #开始时间
    endtime      = peewee.IntegerField(default = int(time.time()))  #结束时间
    createdate   = peewee.IntegerField(default = int(time.time()))  #创建时间

##############################################################################################
class IPTask(BaseModel):
    """ip扫描任务详表"""
    id           = peewee.PrimaryKeyField()
    task_id      = peewee.ForeignKeyField(Task,related_name='iptaskid')          #taskid对应到Task表中的taskid
    task_host    = peewee.CharField()                       #ip地址(127.0.0.1)(192.168.1.1/255)(192.168.1.1-192.168.255.255)
    task_port    = peewee.CharField(default = 80)           #端口 (80)(80,8080,8888,...)(1-65535)
    task_ishttp  = peewee.BooleanField(default = True)      #是否进行banner扫描
    task_isforced= peewee.BooleanField(default = False)     #是否强制扫描
    createdate   = peewee.IntegerField(default = int(time.time()))

class IPResult(BaseModel):
    """ip扫描结果中间表"""
    id           = peewee.PrimaryKeyField()
    ip           = peewee.CharField(unique = True, max_length = 15)
    createdate   = peewee.IntegerField(default = int(time.time()))

class PortResult(BaseModel):
    """端口扫描结果"""
    id           = peewee.PrimaryKeyField()
    host         = peewee.ForeignKeyField(IPResult,related_name='porthost')
    port         = peewee.CharField()
    status_code  = peewee.CharField(default = 0)
    headers      = peewee.TextField(default = '')
    content      = peewee.TextField(default = '')
    banner       = peewee.CharField(default = '') #服务器版本
    title        = peewee.CharField(default = '')
    createdate   = peewee.IntegerField(default = int(time.time()))

class DomainResult(BaseModel):
    """子域名扫描结果"""
    id           = peewee.PrimaryKeyField()
    host         = peewee.ForeignKeyField(IPResult,related_name='domainhost')
    domain       = peewee.CharField(unique = True)
    createdate   = peewee.IntegerField(default = int(time.time()))

class PathResult(BaseModel):
    """目录扫描结果"""
    id           = peewee.PrimaryKeyField()
    domain       = peewee.ForeignKeyField(DomainResult,related_name='pathdomain')
    fid          = peewee.IntegerField(default = 0)  #父目录id 0-为根目录
    isfolder     = peewee.BooleanField(default = True) #是否是文件夹 （文件）
    name         = peewee.CharField()
    path         = peewee.CharField() #全路径

##############################################################################################################
#暂时还不知道怎么写 我先把端口扫描，子域名扫描，目录扫描写完吧
class URLTask(BaseModel):
    id           = peewee.PrimaryKeyField()
    task_id      = peewee.ForeignKeyField(Task,related_name='urltaskid')
    url          = peewee.CharField()
    createdate   = peewee.IntegerField(default = int(time.time()))


##############################################################################################################
if __name__ == '__main__':
    import util

    def createtable(fail_silently=True):
        USERDATA.create_tables([
            Plugin,
            Task,
            IPTask,
            IPResult,
            PortResult,
            DomainResult,
            PathResult])

    def createdata():
        Pg = Plugin.create(plugin_type='portscan',plugin_path='portscan.py',plugin_name='端口扫描')
        Tk = Task.create(task_type=Pg)
        IPTask.create(task_id=Tk,task_host="218.85.1.1-218.85.254.254",task_port='80')
        Ip = IPResult.create(ip='127.0.0.1')
        PortResult.create(host=Ip,port=80)
        Dn = DomainResult.create(host=Ip,domain='localhost')
        PathResult.create(domain=Dn,name='sads',path='./sss')

    def test1():
        for v in PortResult.select().join(IPResult).join(DomainResult).join(PathResult).where(PathResult.name=='sads'):
            print(v.port)

    Ip,c = IPResult.create_or_get(ip='127.0.0.2')
    Dn = DomainResult.create(host=Ip,domain='localhost1')


